// Copyright (c) 2015-2022, Lawrence Livermore National Security, LLC.
// See top-level LICENSE file for details.

/// \file CaliFunctional.h
/// Caliper C++ Functional Annotation Utilities

#ifndef CALI_FUNCTIONAL_H
#define CALI_FUNCTIONAL_H

#define GCC_VERSION (__GNUC__ * 10000 + __GNUC_MINOR__ * 100 + __GNUC_PATCHLEVEL__)
/* Test for GCC > 3.2.0 */

#define VARIADIC_RETURN_SAFE

#include "Annotation.h"
#include "cali_definitions.h"

#include "common/Variant.h"

#include <cstring>
#include <iostream>
#include <tuple>
#include <type_traits>

namespace cali
{
struct SafeAnnotation {
public:

    //ValidatedAnnotation() : inner_annot("ERROR"){}
    SafeAnnotation(const char* name, int opt = 0) : inner_annot(name, opt) { inner_annot = Annotation(name, opt); }

    SafeAnnotation(SafeAnnotation& other) : inner_annot(other.getAnnot()) {}

    Annotation& getAnnot() { return inner_annot; }

    SafeAnnotation& operator= (SafeAnnotation& other)
    {
        inner_annot = other.getAnnot();
        return *this;
    }

    SafeAnnotation& begin()
    {
        inner_annot.begin();
        return *this;
    }

    SafeAnnotation& begin(int data)
    {
        inner_annot.begin(data);
        return *this;
    }

    SafeAnnotation& begin(double data)
    {
        inner_annot.begin(data);
        return *this;
    }

    SafeAnnotation& begin(const char* data)
    {
        inner_annot.begin(Variant(CALI_TYPE_STRING, data, strlen(data)));
        return *this;
    }

    SafeAnnotation& begin(cali_attr_type type, void* data, uint64_t size)
    {
        inner_annot.begin(type, data, size);
        return *this;
    }

    SafeAnnotation& begin(const Variant& data)
    {
        inner_annot.begin(data);
        return *this;
    }

    //set
    SafeAnnotation& set(int data)
    {
        inner_annot.set(data);
        return *this;
    }

    SafeAnnotation& set(double data)
    {
        inner_annot.set(data);
        return *this;
    }

    SafeAnnotation& set(const char* data)
    {
        inner_annot.set(Variant(CALI_TYPE_STRING, data, strlen(data)));
        return *this;
    }

    SafeAnnotation& set(cali_attr_type type, void* data, uint64_t size)
    {
        inner_annot.set(type, data, size);
        return *this;
    }

    SafeAnnotation& set(const Variant& data)
    {
        inner_annot.set(data);
        return *this;
    }

    void end() { inner_annot.end(); }

    template <typename T>
    SafeAnnotation& begin(T start)
    {
        inner_annot.begin("Unmeasurable");
        return *this;
    }

    template <typename T>
    SafeAnnotation& set(T start)
    {
        inner_annot.set("Unmeasurable");
        return *this;
    }

    Annotation inner_annot;
};

cali::SafeAnnotation& wrapper_annotation()
{
    static cali::SafeAnnotation instance("wrapped_function");
    return instance;
}

template <int N>
const char* annotation_name()
{
    static std::string name = "function_argument_" + std::to_string(N);
    return (name.c_str());
}

template <int N>
cali::SafeAnnotation& arg_annotation()
{
    static cali::SafeAnnotation instance(annotation_name<N>());
    return instance;
}

template <int N>
cali::Annotation& arg_annotation_raw()
{
    static cali::Annotation instance(annotation_name<N>());
    return instance;
}

//Wrap a call to a function
template <typename LB, typename... Args>
auto wrap(const char* name, LB body, Args... args) -> typename std::result_of<LB(Args...)>::type
{
    Annotation              startedAnnot = wrapper_annotation().begin(name).getAnnot();
    cali::Annotation::Guard func_annot(startedAnnot);
    return body(args...);
}

cali::Annotation& dummy_annot()
{
    static cali::Annotation instance("wrapped func");
    return instance;
}

template <int N, typename... Args>
struct ArgRecorder {
    static auto record(const char* name) -> std::tuple<> { return std::make_tuple(); }
};

template <int N, typename Arg, typename... Args>
struct ArgRecorder<N, Arg, Args...> {
    using NextRecorder = ArgRecorder<N + 1, Args...>;

    static auto record(const char* name, Arg arg, Args... args) -> decltype(std::tuple_cat(
        NextRecorder::record(name, args...),
        std::move(std::forward_as_tuple(std::move(cali::Annotation::Guard(dummy_annot()))))
    ))
    {
        cali::Annotation::Guard func_annot(arg_annotation<N>().begin(arg).getAnnot());
        return std::tuple_cat(NextRecorder::record(name, args...), std::forward_as_tuple(std::move(func_annot)));
    }
};

//TODO: fix
template <typename LB, typename... Args>
auto wrap_with_args(const char* name, LB body, Args... args) -> typename std::result_of<LB(Args...)>::type
{
    Annotation              startedAnnot = wrapper_annotation().begin(name).getAnnot();
    cali::Annotation::Guard func_annot(startedAnnot);
#ifdef VARIADIC_RETURN_SAFE
    auto n = ArgRecorder<1, Args...>::record(name, args...);
#else
#warning CALIPER WARNING: This  C++ compiler has bugs which prevent argument profiling
#endif
    return body(args...);
}

//Functor containing a function which should always be wrapped. Should not be instantiated directly,
//but through calls to wrap_function below
template <class LB>
struct WrappedFunction {
    WrappedFunction(const char* func_name, LB func) : body(func) { name = func_name; }

    template <typename... Args>
    auto operator() (Args... args) -> typename std::result_of<LB(Args...)>::type
    {

        cali::Annotation::Guard func_annot(wrapper_annotation().begin(name).getAnnot());
        return body(args...);
    }

    LB          body;
    const char* name;
};

template <class LB>
struct ArgWrappedFunction {
    ArgWrappedFunction(const char* func_name, LB func) : body(func) { name = func_name; }

    template <typename... Args>
    auto operator() (Args... args) -> typename std::enable_if<
        !std::is_same<typename std::result_of<LB(Args...)>::type, void>::value,
        typename std::result_of<LB(Args...)>::type>::type
    {
        cali::Annotation::Guard func_annot(wrapper_annotation().begin(name).getAnnot());
#ifdef VARIADIC_RETURN_SAFE
        auto n = ArgRecorder<1, Args...>::record(name, args...);
#else
#warning CALIPER WARNING: This  C++ compiler has bugs which prevent argument profiling
#endif
        auto             return_value = body(args...);
        cali::Annotation return_value_annot("return");
        return_value_annot.set(return_value);
        return_value_annot.end();
        return return_value;
    }

    template <typename... Args>
    auto operator() (Args... args) -> typename std::enable_if<
        std::is_same<typename std::result_of<LB(Args...)>::type, void>::value,
        typename std::result_of<LB(Args...)>::type>::type
    {
        cali::Annotation::Guard func_annot(wrapper_annotation().begin(name).getAnnot());
#ifdef VARIADIC_RETURN_SAFE
        auto n = ArgRecorder<1, Args...>::record(name, args...);
#else
#warning CALIPER WARNING: This  C++ compiler has bugs which prevent argument profiling
#endif
        return body(args...);
    }

    LB          body;
    const char* name;
};

//Helper factory function to create WrappedFunction objects
template <typename LB>
WrappedFunction<LB> wrap_function(const char* name, LB body)
{
    return WrappedFunction<LB>(name, body);
}

template <typename LB>
ArgWrappedFunction<LB> wrap_function_and_args(const char* name, LB body)
{
    return ArgWrappedFunction<LB>(name, body);
}

template <typename InstanceType>
struct Recordable {
    static void record(InstanceType instance) {};
};

template <typename... Args>
struct MapRecordingOperator {
    static auto record() -> void {}
};

template <typename MaybeRecordable>
void recordIfPossible(const MaybeRecordable instance)
{
    Recordable<MaybeRecordable>::record(instance);
}

template <typename MaybeRecordable>
void recordIfPossible(const MaybeRecordable* instance)
{
    Recordable<MaybeRecordable>::record(*instance);
}

cali::Annotation argument_number_annot("argument");

template <typename Arg, typename... Args>
struct MapRecordingOperator<Arg, Args...> {
    static auto record(Arg arg, Args... args) -> void
    {
        argument_number_annot.set(static_cast<int>(sizeof...(args)));
        recordIfPossible(arg);
        argument_number_annot.end(); // can be moved to callsite for speed, looks bad though
        MapRecordingOperator<Args...>::record(args...);
    }
};

template <typename... Args>
void recordAll(Args... args)
{
    MapRecordingOperator<Args...>::record(args...);
}

cali::Annotation recording_phase("recording_phase");

template <class LB>
struct RecordedFunction {
    RecordedFunction(const char* func_name, LB func) : body(func) { name = func_name; }

    template <typename... Args>
    auto operator() (Args... args) -> typename std::enable_if<
        !std::is_same<typename std::result_of<LB(Args...)>::type, void>::value,
        typename std::result_of<LB(Args...)>::type>::type
    {
        recording_phase.begin("pre");
        recordAll(args...);
        recording_phase.end();
        cali::Annotation::Guard func_annot(cali::Annotation("debugged_function").begin(name));
        auto                    return_value = body(args...);
        //cali::Annotation return_value_annot("return");
        //return_value_annot.set(return_value);
        //return_value_annot.end();
        recording_phase.begin("post");
        recordAll(args...);
        recording_phase.end();
        return return_value;
    }

    template <typename... Args>
    auto operator() (Args... args) -> typename std::enable_if<
        std::is_same<typename std::result_of<LB(Args...)>::type, void>::value,
        typename std::result_of<LB(Args...)>::type>::type
    {
        cali::Annotation::Guard func_annot(cali::Annotation("debugged_function").begin(name));
        recording_phase.begin("pre");
        recordAll(args...);
        recording_phase.end();
        body(args...);
        recording_phase.begin("post");
        recordAll(args...);
        recording_phase.end();
    }

    LB          body;
    const char* name;
};

template <typename LB>
RecordedFunction<LB> make_recorded_function(const char* name, LB body)
{
    return RecordedFunction<LB>(name, body);
}

} // end namespace cali
#endif //CALI_ANNOTATED_REGION_H
